
package components.tree;

import java.util.ArrayList;
import java.util.Stack;

import javax.swing.tree.DefaultTreeSelectionModel;
import javax.swing.tree.TreeModel;
import javax.swing.tree.TreePath;
import javax.swing.tree.TreeSelectionModel;

public class CheckTreeSelectionModel extends DefaultTreeSelectionModel{ 
  private TreeModel model; 

  public CheckTreeSelectionModel(TreeModel model){ 
      this.model = model; 
      setSelectionMode(TreeSelectionModel.DISCONTIGUOUS_TREE_SELECTION); 
  } 

  // tests whether there is any unselected node in the subtree of given path 
  public boolean isPartiallySelected(TreePath path){ 
      if(isPathSelected(path, true)) 
          return false; 
      TreePath[] selectionPaths = getSelectionPaths(); 
      if(selectionPaths==null) 
          return false; 
      for(int j = 0; j<selectionPaths.length; j++){ 
          if(isDescendant(selectionPaths[j], path)) 
              return true; 
      } 
      return false; 
  } 

  // tells whether given path is selected. 
  // if dig is true, then a path is assumed to be selected, if 
  // one of its ancestor is selected. 
  public boolean isPathSelected(TreePath path, boolean dig){ 
      if(!dig) 
          return super.isPathSelected(path); 
      while(path!=null && !super.isPathSelected(path)) 
          path = path.getParentPath(); 
      return path!=null; 
  } 

  // is path1 descendant of path2 
  private boolean isDescendant(TreePath path1, TreePath path2){ 
      Object obj1[] = path1.getPath(); 
      Object obj2[] = path2.getPath(); 
      for(int i = 0; i<obj2.length; i++){ 
          if(obj1[i]!=obj2[i]) 
              return false; 
      } 
      return true; 
  } 

  public void setSelectionPaths(TreePath[] pPaths){ 
      throw new UnsupportedOperationException("not implemented yet!!!"); 
  } 

  public void addSelectionPaths(TreePath[] paths){ 
      // unselect all descendants of paths[] 
      for(int i = 0; i<paths.length; i++){ 
          TreePath path = paths[i]; 
          TreePath[] selectionPaths = getSelectionPaths(); 
          if(selectionPaths==null) 
              break; 
          ArrayList toBeRemoved = new ArrayList(); 
          for(int j = 0; j<selectionPaths.length; j++){ 
              if(isDescendant(selectionPaths[j], path)) 
                  toBeRemoved.add(selectionPaths[j]); 
          } 
          super.removeSelectionPaths((TreePath[])toBeRemoved.toArray(new TreePath[0])); 
      } 

      // if all siblings are selected then unselect them and select parent recursively 
      // otherwize just select that path. 
      for(int i = 0; i<paths.length; i++){ 
          TreePath path = paths[i]; 
          TreePath temp = null; 
          while(areSiblingsSelected(path)){ 
              temp = path; 
              if(path.getParentPath()==null) 
                  break; 
              path = path.getParentPath(); 
          } 
          if(temp!=null){ 
              if(temp.getParentPath()!=null) 
                  addSelectionPath(temp.getParentPath()); 
              else{ 
                  if(!isSelectionEmpty()) 
                      removeSelectionPaths(getSelectionPaths()); 
                  super.addSelectionPaths(new TreePath[]{temp}); 
              } 
          }else 
              super.addSelectionPaths(new TreePath[]{ path}); 
      } 
  } 

  // tells whether all siblings of given path are selected. 
  private boolean areSiblingsSelected(TreePath path){ 
      TreePath parent = path.getParentPath(); 
      if(parent==null) 
          return true; 
      Object node = path.getLastPathComponent(); 
      Object parentNode = parent.getLastPathComponent(); 

      int childCount = model.getChildCount(parentNode); 
      for(int i = 0; i<childCount; i++){ 
          Object childNode = model.getChild(parentNode, i); 
          if(childNode==node) 
              continue; 
          if(!isPathSelected(parent.pathByAddingChild(childNode))) 
              return false; 
      } 
      return true; 
  } 

  public void removeSelectionPaths(TreePath[] paths){ 
      for(int i = 0; i<paths.length; i++){ 
          TreePath path = paths[i]; 
          if(path.getPathCount()==1) 
              super.removeSelectionPaths(new TreePath[]{ path}); 
          else 
              toggleRemoveSelection(path); 
      } 
  } 

  // if any ancestor node of given path is selected then unselect it 
  //  and selection all its descendants except given path and descendants. 
  // otherwise just unselect the given path 
  private void toggleRemoveSelection(TreePath path){ 
      Stack stack = new Stack(); 
      TreePath parent = path.getParentPath(); 
      while(parent!=null && !isPathSelected(parent)){ 
          stack.push(parent); 
          parent = parent.getParentPath(); 
      } 
      if(parent!=null) 
          stack.push(parent); 
      else{ 
          super.removeSelectionPaths(new TreePath[]{path}); 
          return; 
      } 

      while(!stack.isEmpty()){ 
          TreePath temp = (TreePath)stack.pop(); 
          TreePath peekPath = stack.isEmpty() ? path : (TreePath)stack.peek(); 
          Object node = temp.getLastPathComponent(); 
          Object peekNode = peekPath.getLastPathComponent(); 
          int childCount = model.getChildCount(node); 
          for(int i = 0; i<childCount; i++){ 
              Object childNode = model.getChild(node, i); 
              if(childNode!=peekNode) 
                  super.addSelectionPaths(new TreePath[]{temp.pathByAddingChild(childNode)}); 
          } 
      } 
      super.removeSelectionPaths(new TreePath[]{parent}); 
  } 
}